import glfw
import glm
from OpenGL.GL import *
from OpenGL.GLU import *
import numpy as np
import trimesh
import time
import ctypes
import open3d as o3d

from PIL import Image
from scipy.spatial.transform import Rotation as R
from learnopengl.ibl.shader import Shader
from learnopengl.ibl.camera import Camera

from utils import load_scan

SCR_WIDTH = 800
SCR_HEIGHT = 600
camera = Camera(glm.vec3(0, 0, 3))
lastX = SCR_WIDTH / 2
lastY = SCR_HEIGHT / 2
firstMouse = True

deltaTime = 0.0
lastFrame = 0.0

euler_angle = [90, 0, 0]
r = R.from_euler('xyz', euler_angle, degrees=True)
rotation_matrix = r.as_matrix()
rotation_matrix_inv = np.linalg.inv(rotation_matrix)

verticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    layout (location = 1) in vec2 aTexCoord;
    
    out vec2 TexCoord;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        gl_Position = projection * view * model * vec4(aPos, 1.0);
        TexCoord = aTexCoord;
    }
'''

fragmentShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    in vec2 TexCoord;

    uniform sampler2D ourTexture;
    
    void main()
    {
        FragColor = texture(ourTexture, TexCoord);
    }
'''

smpl_verticesShaderSource = '''
    #version 330 core
    layout (location = 0) in vec3 aPos;
    
    uniform mat4 model;
    uniform mat4 view;
    uniform mat4 projection;
    
    void main()
    {
        gl_Position = projection * view * model * vec4(aPos, 1.0);
    }
'''

smpl_fragmentShaderSource = '''
    #version 330 core
    out vec4 FragColor;
    
    void main()
    {
        FragColor = vec4(1.0f, 0.5f, 0.4f, 1.0f);
    }
'''

def load_texture(texture_path, texture_unit, flip=True):
    img = Image.open(texture_path).convert('RGB')
    if flip:
        img = img.transpose(Image.FLIP_TOP_BOTTOM)
    img_data = img.tobytes()
    width, height = img.size

    texture = glGenTextures(1)
    glActiveTexture(GL_TEXTURE0 + texture_unit)
    glBindTexture(GL_TEXTURE_2D, texture)

    glTexImage2D(GL_TEXTURE_2D, 0, GL_RGB, width, height, 0, GL_RGB, GL_UNSIGNED_BYTE, img_data)
    glGenerateMipmap(GL_TEXTURE_2D)

    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAX_LEVEL, 3)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_S, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_WRAP_T, GL_CLAMP_TO_EDGE)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MAG_FILTER, GL_LINEAR)
    glTexParameteri(GL_TEXTURE_2D, GL_TEXTURE_MIN_FILTER, GL_LINEAR_MIPMAP_LINEAR)


    return texture

def framebuffer_size_callback(window, width, height):
    glViewport(0, 0, width, height)

def processInput(window):
    global camera, deltaTime
    if glfw.get_key(window, glfw.KEY_W) == glfw.PRESS:
        camera.ProcessKeyboard("FORWARD", deltaTime)
    if glfw.get_key(window, glfw.KEY_S) == glfw.PRESS:
        camera.ProcessKeyboard("BACKWARD", deltaTime)
    if glfw.get_key(window, glfw.KEY_A) == glfw.PRESS:
        camera.ProcessKeyboard("LEFT", deltaTime)
    if glfw.get_key(window, glfw.KEY_D) == glfw.PRESS:
        camera.ProcessKeyboard("RIGHT", deltaTime)
    if glfw.get_key(window, glfw.KEY_Q) == glfw.PRESS:
        camera.ProcessKeyboard("UP", deltaTime)
    if glfw.get_key(window, glfw.KEY_E) == glfw.PRESS:
        camera.ProcessKeyboard("DOWN", deltaTime)
    if glfw.get_key(window, glfw.KEY_LEFT) == glfw.PRESS:
        camera.ProcessKeyboard("ROTATE_Positive", deltaTime)   
    if glfw.get_key(window, glfw.KEY_RIGHT) == glfw.PRESS:
        camera.ProcessKeyboard("ROTATE_Negative", deltaTime) 
    
    if glfw.get_key(window, glfw.KEY_ESCAPE) == glfw.PRESS:
        glfw.set_window_should_close(window, True)

def scroll_callback(window, xoffset, yoffset):
    global camera
    camera.ProcessMouseScroll(yoffset)

# 将 3x3 旋转矩阵转换为 4x4 齐次矩阵
def to_homogeneous(matrix):
    assert matrix.shape == (3, 3)
    homogeneous_matrix = np.eye(4)
    homogeneous_matrix[:3, :3] = matrix
    return homogeneous_matrix

def keep_largest(mesh):
    mesh_lst = mesh.split(only_watertight=False)
    keep_mesh = mesh_lst[0]
    for mesh in mesh_lst:
        if mesh.vertices.shape[0] > keep_mesh.vertices.shape[0]:
            keep_mesh = mesh
    return keep_mesh



if __name__ == "__main__":
    file_idx = "0013"
    obj_path = f'THuman2.1_Release/model/{file_idx}/{file_idx}.obj'
    # smpl_obj_path = f'THuman2.1_Release/smplx/{file_idx}/mesh_smplx.obj'
    # spml_path = f"THuman2.1_Release/smplx/{file_idx}/smplx_param.pkl"
    texture_path = f'THuman2.1_Release/model/{file_idx}/material0.jpeg'
    # 加载网格和纹理
    mesh = trimesh.load(obj_path, process=False)
    vertices_ori = np.array(mesh.vertices, dtype=np.float32)
    faces = np.array(mesh.faces.flatten(), dtype=np.int32)
    vertices_ori, faces = remove_isolated_vertices(vertices_ori, faces)
    # vertices, faces, normals, faces_normals, textures, face_textures = load_scan(obj_path, with_normal=True, with_texture=True)
    # vertices_ori = np.array(vertices, dtype=np.float32)
    # faces = np.array(faces.flatten(), dtype=np.uint32)
    
    # 计算变换矩阵
    # vertices_ori = np.matmul(rotation_matrix_inv, vertices_ori.T).T
    min_vals = vertices_ori.min(axis=0)
    max_vals = vertices_ori.max(axis=0)
    center = (max_vals + min_vals) / 2.0
    scale = (max_vals - min_vals).max() / 2.0
    vertices_ori = (vertices_ori - center) / scale

    
    # try:
    #     # uv_coords = np.array(textures, dtype=np.float32)
    #     uv_coords = np.array(mesh.visual.uv, dtype=np.float32)
    #     vertices = np.concatenate([vertices_ori, uv_coords], axis=1)
    # except:
    #     uv_coords = None
    #     vertices = vertices_ori
    uv_coords = None
    vertices = vertices_ori

    # 初始化
    if not glfw.init():
        raise Exception("GLFW can't be initialized")
    glfw.window_hint(glfw.CONTEXT_VERSION_MAJOR, 3)
    glfw.window_hint(glfw.CONTEXT_VERSION_MINOR, 3)
    glfw.window_hint(glfw.OPENGL_PROFILE, glfw.OPENGL_CORE_PROFILE)
    
    window = glfw.create_window(SCR_WIDTH, SCR_HEIGHT, "LearnOpenGL", None, None)
    if not window:
        glfw.terminate()
        raise Exception("GLFW window can't be created")

    glfw.make_context_current(window)
    glfw.set_framebuffer_size_callback(window, framebuffer_size_callback)
    glfw.set_scroll_callback(window, scroll_callback)
    
    glfw.set_input_mode(window, glfw.CURSOR, glfw.CURSOR_DISABLED)
    
    glEnable(GL_DEPTH_TEST)
    
    # 打印 GPU 信息
    print("Renderer:", glGetString(GL_RENDERER).decode())
    print("OpenGL version supported:", glGetString(GL_VERSION).decode())
    
    VAO = glGenVertexArrays(1)
    VBO = glGenBuffers(1)
    EBO = glGenBuffers(1)
    
    glBindVertexArray(VAO)
    
    glBindBuffer(GL_ARRAY_BUFFER, VBO)
    glBufferData(GL_ARRAY_BUFFER, vertices.nbytes, vertices, GL_STATIC_DRAW)
    
    glBindBuffer(GL_ELEMENT_ARRAY_BUFFER, EBO)
    glBufferData(GL_ELEMENT_ARRAY_BUFFER, faces.nbytes, faces, GL_STATIC_DRAW)

    has_texture = True
    if uv_coords is not None:
        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, vertices.itemsize * 5, ctypes.c_void_p(0))
        glEnableVertexAttribArray(0)
        glVertexAttribPointer(1, 2, GL_FLOAT, GL_FALSE, vertices.itemsize * 5, ctypes.c_void_p(3 * vertices.itemsize))
        glEnableVertexAttribArray(1)
    else:
        glVertexAttribPointer(0, 3, GL_FLOAT, GL_FALSE, vertices.itemsize * 3, ctypes.c_void_p(0))
        glEnableVertexAttribArray(0)
        has_texture = False
    
    if has_texture:
        texture_map = load_texture(texture_path, 0) 
    
    if has_texture:
        shaderProgram = Shader(verticesShaderSource, fragmentShaderSource)
    else:
        shaderProgram = Shader(smpl_verticesShaderSource, smpl_fragmentShaderSource)
    shaderProgram.use()
    glUniform1i(glGetUniformLocation(shaderProgram.ID, "ourTexture"), 0)

    
    rotation_angle = 0.0
    while not glfw.window_should_close(window):  
        currentFrame = glfw.get_time()
        deltaTime = currentFrame - lastFrame
        lastFrame = currentFrame
      
        processInput(window)

        glClearColor(0.2, 0.3, 0.3, 1.0)
        glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
        
        shaderProgram.use()
        project = glm.perspective(glm.radians(camera.zoom), SCR_WIDTH / SCR_HEIGHT, 0.1, 100.0)
        # project = glm.mat4(np.dot(np.array(project), rotation_matrix_inv))
        view = camera.GetViewMatrix()
        shaderProgram.set_mat4("projection", project)
        shaderProgram.set_mat4("view", view)
        
        model = glm.rotate(glm.mat4(1.0), glm.radians(camera.yaw), glm.vec3(0.0, 1.0, 0.0))
        # model_xd = glm.rotate(glm.mat4(1.0), glm.radians(-90), glm.vec3(1.0, 0.0, 0.0))
        # model = model @ model_xd
        model = model @ glm.mat4(to_homogeneous(rotation_matrix_inv))

        shaderProgram.set_mat4("model", model)
        
        
        if has_texture:
            glActiveTexture(GL_TEXTURE0)
            glBindTexture(GL_TEXTURE_2D, texture_map)
        
        glBindVertexArray(VAO)
        glDrawElements(GL_TRIANGLES, len(faces), GL_UNSIGNED_INT, None)

        glfw.swap_buffers(window)
        glfw.poll_events()

    glfw.terminate()
